package proxy

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"

	"ehang.io/nps/lib/common"
	"ehang.io/nps/lib/conn"
	"ehang.io/nps/lib/file"
	"github.com/astaxie/beego/logs"
)

// HTTPDisguiseServer HTTP伪装的TCP隧道服务器
type HTTPDisguiseServer struct {
	BaseServer
	listener net.Listener
}

// NewHTTPDisguiseServer 创建HTTP伪装TCP隧道服务器
func NewHTTPDisguiseServer(bridge NetBridge, task *file.Tunnel) *HTTPDisguiseServer {
	return &HTTPDisguiseServer{
		BaseServer: BaseServer{
			bridge: bridge,
			task:   task,
			Mutex:  sync.Mutex{},
		},
	}
}

// Start 启动HTTP伪装TCP隧道服务
func (s *HTTPDisguiseServer) Start() error {
	return conn.NewTcpListenerAndProcess(s.task.ServerIp+":"+strconv.Itoa(s.task.Port), func(c net.Conn) {
		if err := s.CheckFlowAndConnNum(s.task.Client); err != nil {
			logs.Warn("client id %d, task id %d, error %s, when http disguise tcp connection", s.task.Client.Id, s.task.Id, err.Error())
			c.Close()
			return
		}
		logs.Trace("new http disguise tcp connection, local port %d, client %d, remote address %s", s.task.Port, s.task.Client.Id, c.RemoteAddr())
		s.processHTTPDisguise(conn.NewConn(c))
		s.task.Client.AddConn()
	}, &s.listener)
}

// Close 关闭服务
func (s *HTTPDisguiseServer) Close() error {
	if s.listener != nil {
		return s.listener.Close()
	}
	return nil
}

// processHTTPDisguise 处理HTTP伪装的TCP连接
func (s *HTTPDisguiseServer) processHTTPDisguise(c *conn.Conn) error {
	defer c.Close()

	// 读取第一个数据包，判断是否为HTTP请求
	buf := make([]byte, 4096)
	n, err := c.Read(buf)
	if err != nil {
		logs.Error("read first packet error:", err)
		return err
	}

	data := buf[:n]
	
	// 检查是否为HTTP请求
	if s.isHTTPRequest(data) {
		// 如果是HTTP请求，按HTTP协议处理
		return s.handleHTTPRequest(c, data)
	} else {
		// 如果不是HTTP请求，按TCP隧道处理，但需要伪装HTTP响应
		return s.handleTCPWithHTTPDisguise(c, data)
	}
}

// isHTTPRequest 检查数据是否为HTTP请求
func (s *HTTPDisguiseServer) isHTTPRequest(data []byte) bool {
	dataStr := string(data)
	httpMethods := []string{"GET", "POST", "PUT", "DELETE", "HEAD", "OPTIONS", "PATCH", "CONNECT"}
	
	for _, method := range httpMethods {
		if strings.HasPrefix(dataStr, method+" ") {
			return true
		}
	}
	return false
}

// handleHTTPRequest 处理真正的HTTP请求
func (s *HTTPDisguiseServer) handleHTTPRequest(c *conn.Conn, firstData []byte) error {
	// 发送伪装的HTTP响应
	response := s.generateFakeHTTPResponse()
	c.Write([]byte(response))
	return nil
}

// handleTCPWithHTTPDisguise 处理TCP数据但伪装成HTTP
func (s *HTTPDisguiseServer) handleTCPWithHTTPDisguise(c *conn.Conn, firstData []byte) error {
	// 获取目标地址
	targetAddr, err := s.task.Target.GetRandomTarget()
	if err != nil {
		logs.Warn("tcp port %d, client id %d, task id %d connect error %s", s.task.Port, s.task.Client.Id, s.task.Id, err.Error())
		return err
	}

	// 创建到目标的连接
	link := conn.NewLink(common.CONN_TCP, targetAddr, s.task.Client.Cnf.Crypt, s.task.Client.Cnf.Compress, c.Conn.RemoteAddr().String(), s.task.Target.LocalProxy)
	target, err := s.bridge.SendLinkInfo(s.task.Client.Id, link, s.task)
	if err != nil {
		logs.Warn("get connection from client id %d error %s", s.task.Client.Id, err.Error())
		return err
	}
	defer target.Close()

	// 创建HTTP伪装的数据传输通道
	return s.startHTTPDisguiseDataTransfer(c, target, firstData)
}

// startHTTPDisguiseDataTransfer 开始HTTP伪装的数据传输
func (s *HTTPDisguiseServer) startHTTPDisguiseDataTransfer(clientConn *conn.Conn, targetConn net.Conn, firstData []byte) error {
	// 发送伪装的HTTP响应头，表示建立了持久连接
	httpResponse := "HTTP/1.1 200 OK\r\n" +
		"Content-Type: application/octet-stream\r\n" +
		"Transfer-Encoding: chunked\r\n" +
		"Connection: keep-alive\r\n" +
		"Cache-Control: no-cache\r\n" +
		"\r\n"
	
	clientConn.Write([]byte(httpResponse))

	// 将第一个数据包发送到目标
	if len(firstData) > 0 {
		_, err := targetConn.Write(firstData)
		if err != nil {
			return err
		}
	}

	// 启动双向数据传输
	var wg sync.WaitGroup
	wg.Add(2)

	// 客户端到目标的数据传输（伪装成HTTP chunk）
	go func() {
		defer wg.Done()
		s.transferClientToTarget(clientConn, targetConn)
	}()

	// 目标到客户端的数据传输（伪装成HTTP chunk）
	go func() {
		defer wg.Done()
		s.transferTargetToClient(targetConn, clientConn)
	}()

	wg.Wait()
	return nil
}

// transferClientToTarget 客户端到目标的数据传输
func (s *HTTPDisguiseServer) transferClientToTarget(clientConn *conn.Conn, targetConn net.Conn) {
	defer targetConn.Close()
	
	buf := make([]byte, 32*1024)
	for {
		n, err := clientConn.Read(buf)
		if err != nil {
			if err != io.EOF {
				logs.Trace("client to target read error:", err)
			}
			break
		}

		// 解析HTTP chunk格式的数据
		data := s.parseHTTPChunk(buf[:n])
		if len(data) > 0 {
			_, err = targetConn.Write(data)
			if err != nil {
				logs.Trace("client to target write error:", err)
				break
			}
			s.task.Flow.Add(int64(len(data)), 0)
		}
	}
}

// transferTargetToClient 目标到客户端的数据传输
func (s *HTTPDisguiseServer) transferTargetToClient(targetConn net.Conn, clientConn *conn.Conn) {
	defer clientConn.Close()
	
	buf := make([]byte, 32*1024)
	for {
		n, err := targetConn.Read(buf)
		if err != nil {
			if err != io.EOF {
				logs.Trace("target to client read error:", err)
			}
			break
		}

		// 将数据包装成HTTP chunk格式
		chunkData := s.wrapAsHTTPChunk(buf[:n])
		_, err = clientConn.Write(chunkData)
		if err != nil {
			logs.Trace("target to client write error:", err)
			break
		}
		s.task.Flow.Add(0, int64(len(buf[:n])))
	}

	// 发送结束chunk
	clientConn.Write([]byte("0\r\n\r\n"))
}

// wrapAsHTTPChunk 将数据包装成HTTP chunk格式
func (s *HTTPDisguiseServer) wrapAsHTTPChunk(data []byte) []byte {
	if len(data) == 0 {
		return []byte("0\r\n\r\n")
	}
	
	chunkSize := fmt.Sprintf("%x\r\n", len(data))
	chunk := make([]byte, 0, len(chunkSize)+len(data)+2)
	chunk = append(chunk, []byte(chunkSize)...)
	chunk = append(chunk, data...)
	chunk = append(chunk, []byte("\r\n")...)
	
	return chunk
}

// parseHTTPChunk 解析HTTP chunk格式的数据
func (s *HTTPDisguiseServer) parseHTTPChunk(data []byte) []byte {
	// 简单的chunk解析，实际应用中可能需要更复杂的状态机
	reader := bufio.NewReader(bytes.NewReader(data))
	
	var result []byte
	for {
		// 读取chunk大小
		line, err := reader.ReadString('\n')
		if err != nil {
			break
		}
		
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		
		// 解析chunk大小
		var chunkSize int64
		if _, err := fmt.Sscanf(line, "%x", &chunkSize); err != nil {
			break
		}
		
		if chunkSize == 0 {
			break
		}
		
		// 读取chunk数据
		chunkData := make([]byte, chunkSize)
		n, err := reader.Read(chunkData)
		if err != nil || int64(n) != chunkSize {
			break
		}
		
		result = append(result, chunkData...)
		
		// 跳过CRLF
		reader.ReadString('\n')
	}
	
	return result
}

// generateFakeHTTPResponse 生成伪装的HTTP响应
func (s *HTTPDisguiseServer) generateFakeHTTPResponse() string {
	return "HTTP/1.1 200 OK\r\n" +
		"Content-Type: text/html; charset=utf-8\r\n" +
		"Content-Length: 100\r\n" +
		"Server: nginx/1.18.0\r\n" +
		"Date: " + time.Now().Format(time.RFC1123) + "\r\n" +
		"Connection: close\r\n" +
		"\r\n" +
		"<html><head><title>Welcome</title></head><body><h1>Server is running normally</h1></body></html>"
}